from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field, model_validator
from langchain_core.embeddings import Embeddings
from llama_index.core.embeddings import BaseEmbedding
from langchain_core.utils import get_from_dict_or_env
import requests


class CustomAPIEmbeddings(BaseModel, Embeddings):
    """
    自定义嵌入模型集成类，可将嵌入请求发送到指定 URL。

    使用示例：

        embed = CustomAPIEmbeddings(
            model="your-model-name",
            embedding_api_url="https://your-embedding-api-endpoint",
            api_key="your-api-key",
        )

        # 单条文本向量
        embedding = embed.embed_query("hello world")

        # 多条文本向量
        embeddings = embed.embed_documents(["hello", "world"])
    """

    model_name: str = Field(default="bge-large-zh")
    api_key: str
    api_url: str  # 你的目标 API URL
    dimensions: Optional[int] = None

    @model_validator(mode="before")
    @classmethod
    def validate_environment(cls, values: Dict) -> Any:
        """支持从环境变量中读取 API_KEY。"""
        values["api_key"] = get_from_dict_or_env(values, "api_key", "MY_API_KEY")
        values["api_url"] = get_from_dict_or_env(values, "api_url", "MY_MODEL_API_URL")
        return values

    def _post_embed_request(self, texts: List[str]) -> List[List[float]]:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"{self.api_key}",
        }
        payload = {
            "model": self.model_name,
            "input": texts,
        }
        if self.dimensions is not None:
            payload["dimensions"] = self.dimensions
        response = requests.post(self.api_url, headers=headers, json=payload)
        if response.status_code != 200:
            raise ValueError(f"Embedding request failed with status {response.status_code}: {response.text}")

        result = response.json()
        # 假设返回格式为 { "data": [ { "embedding": [...] }, ... ] }
        return [item["embedding"] for item in result.get("data", [])]

    def embed_query(self, text: str) -> List[float]:
        return self._post_embed_request([text])[0]

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return self._post_embed_request(texts)

class InstructorEmbeddings(BaseEmbedding):
    def __init__(
        self,
        embedding_client: Embeddings,
        instruction: str = "Represent the Computer Science documentation or question:",
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self._model = embedding_client  # ✅ 用你的自定义 embedding 实例
        self._instruction = instruction

    def _get_query_embedding(self, query: str) -> List[float]:
        full_input = f"{self._instruction} {query}"
        return self._model.embed_query(full_input)

    def _get_text_embedding(self, text: str) -> List[float]:
        full_input = f"{self._instruction} {text}"
        return self._model.embed_query(full_input)  # 也可用 embed_documents([...])[0]

    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        full_inputs = [f"{self._instruction} {t}" for t in texts]
        return self._model.embed_documents(full_inputs)

    async def _aget_query_embedding(self, query: str) -> List[float]:
        return self._get_query_embedding(query)

    async def _aget_text_embedding(self, text: str) -> List[float]:
        return self._get_text_embedding(text)